from abc import ABC, abstractclassmethod

import numpy as np
from tqdm import trange

from .tensor import Tensor


def get_parameters(obj):
    parameters = []
    if isinstance(obj, Tensor):
        parameters.append(obj)
    elif isinstance(obj, list):
        for x in obj:
            parameters.extend(get_parameters(x))
    elif hasattr(obj, '__dict__'):
        for k, v in obj.__dict__.items():
            parameters.extend(get_parameters(v))
    return parameters


class Model(ABC):
    def __init__(self) -> None:
        super().__init__()

    @abstractclassmethod
    def forward(self, x):
        ...

    def get_params(self):
        return get_parameters(self)

    def train(self, X_train, y_train, optim, epochs=10, BS=128, lossfn=lambda x: x.mean()):
        Tensor.training = True
        losses, accuracies = [], []
        with trange(epochs) as t:
            for _ in t:
                for _ in range(len(y_train)//BS+1):
                    samp = np.random.randint(0, X_train.shape[0], size=BS)

                    x = Tensor(X_train[samp])
                    y = y_train[samp]

                    # network
                    out = self.forward(x)

                    loss = lossfn(out, y)
                    optim.zero_grad()
                    loss.backward()
                    optim.step()

                    cat = np.argmax(out.data, axis=-1)
                    accuracy = (cat == y).mean()

                    # printing
                    loss = loss.data[0]
                    losses.append(loss)
                    accuracies.append(accuracy)
                    t.set_description("loss %.5f accuracy %.5f" %
                                      (loss, accuracy*100))
        return losses, accuracies

    def evaluate(self, X_test, y_test, BS=256):
        Tensor.training = False
        sm = 0
        with trange((len(y_test)-1)//BS+1) as t:
            for i in t:
                if i != 0:
                    t.set_description('Acc: {:.3f}%'.format(sm/(i*BS)*100))
                out = self.forward(
                    Tensor(X_test[i*BS:(i+1)*BS])).data.argmax(-1)
                sm += (out == y_test[i*BS:(i+1)*BS]).sum()
        return sm/len(y_test)

    def predict(self, X, BS=256):
        Tensor.training = False
        N = len(X)
        outs = []
        with trange((N-1)//BS+1) as t:
            for i in t:
                outs.append(self.forward(
                    Tensor(X[i*BS:(i+1)*BS])).data)
        return np.concatenate(outs, axis=0)
